package inbox

import (
	"context"
	"fmt"
	"io"
	"log"

	"agent.com/iagentgo/facade"
	"agent.com/iagentgo/model"
	"github.com/cloudwego/hertz/pkg/app"
	"github.com/cloudwego/hertz/pkg/app/server"
	"github.com/cloudwego/hertz/pkg/protocol/sse"
	"github.com/hertz-contrib/websocket"
)

var upgrader = websocket.HertzUpgrader{}

type wsOut struct {
	facade.ChatSendOut
	conn *websocket.Conn
}

type sseOut struct {
	facade.ChatSendOut
	writer *sse.Writer
}

func (w *wsOut) Send(req *facade.ChatResponse) error {
	err := w.conn.WriteMessage(websocket.TextMessage, []byte(req.Answer))
	return err
}

func (s *sseOut) Send(req *facade.ChatResponse) error {
	err := s.writer.WriteEvent("id-x", "message", []byte(req.Answer))
	return err
}

func hello(ctx context.Context, c *app.RequestContext) {
	c.String(200, "Hello, World!")
}

func wsHandler(conn *websocket.Conn) {
	uid := conn.RemoteAddr().String()
	log.Printf("ws connected, remote addr: %s", uid)
	for {
		mt, message, err := conn.ReadMessage()
		if err != nil {
			break
		}
		if mt == websocket.PingMessage || mt == websocket.PongMessage {
			conn.WriteMessage(mt, message)
			continue
		}
		req := &facade.ChatRequest{
			UserId:   uid,
			Role:     facade.RoleChat,
			ModelId:  facade.QWEN,
			Question: string(message),
			Data:     "",
		}

		bg := model.GetModelBridge()
		stream, err := bg.ChatStream(facade.ModelStyleDefault, req)
		if err != nil {
			log.Println(err)
			body := fmt.Sprintf("error: %v", err)
			conn.WriteMessage(websocket.TextMessage, []byte(body))
			continue
		}
		out := &wsOut{
			conn: conn,
		}
		err = stream.Recv(out)
		if (err != nil) && (err != io.EOF) {
			log.Println(err)
			body := fmt.Sprintf("error: %v", err)
			conn.WriteMessage(websocket.TextMessage, []byte(body))
			continue
		}
	}
	log.Printf("ws closed, remote addr: %s", uid)
	conn.Close()
}

func wsChat(ctx context.Context, c *app.RequestContext) {
	err := upgrader.Upgrade(c, wsHandler)
	if err != nil {
		c.String(500, err.Error())
		return
	}
}

func sseChat(ctx context.Context, c *app.RequestContext) {
	req := &facade.ChatRequest{}
	err := c.Bind(req)
	if err != nil {
		c.String(500, err.Error())
		return
	}
	bg := model.GetModelBridge()
	stream, err := bg.ChatStream(facade.ModelStyleDefault, req)
	if err != nil {
		c.String(500, err.Error())
		return
	}
	w := sse.NewWriter(c)
	sseOut := &sseOut{
		writer: w,
	}
	err = stream.Recv(sseOut)
	if (err != nil) && (err != io.EOF) {
		c.String(500, err.Error())
		return
	}

	w.Close()
}

func AddHertzRoute(h *server.Hertz) {
	h.GET("/", hello)
	h.GET("/ws", wsChat)
	h.GET("/sse", sseChat)
}
